from tqdm import tqdm
from preprossData.utils.crop_padding import crop_resize_save
import re
import csv
import json
import sys
import os
sys.path.insert(0, os.path.abspath(
    os.path.join(os.path.dirname(__file__), '..', '..')))

# ---------------------------
# 诊断代码映射表
diagnosis_mapping = {
    "0": "Physiological finding",
    "1": "Retinopathy of Prematurity Stage 0",
    "2": "Retinopathy of Prematurity Stage 1",
    "3": "Retinopathy of Prematurity Stage 2",
    "4": "Retinopathy of Prematurity Stage 3",
    "5": "Retinopathy of Prematurity Stage 4 A",
    "6": "Retinopathy of Prematurity Stage 4B",
    "7": "Retinopathy of Prematurity Stage 5",
    "8": "Aggressive Retinopathy of Prematurity",
    "9": "Status post Retinopathy of Prematurity",
    "10": "Retinal astrocytic hamartomas",
    "11": "Retinal hemorrhages - abnormal bleeding within the delicate retinal blood vessels",
    "12": "Optic Nerve Hypoplasia",
    "13": "Toxoplasmosis chorioretinitis - infection caused by the Toxoplasma gondii parasite"
}

# Plus Form 映射：0 表示无，1 表示 pre-plus，2 表示 plus
plus_mapping = {
    "0": "",
    "1": "retinopathy of prematurity pre-plus form",
    "2": "retinopathy of prematurity plus form"
}


def parse_info(csv_path):
    """
    读取 CSV 文件，并返回一个字典，格式如下：
    {
        "<ID>_<SERIE NUMBER>": {
            "patient info": { ... },  # 包含其它字段（字段名已清洗为小写，去除括号内内容）
            "text": "<diagnosis text>[, <plus form text>]"
        },
        ...
    }
    排除的字段：device, plus form, diagnosis code, id, serie number
    对 SEX 字段进行处理：'F' 替换为 'Female'，'M' 替换为 'Male'
    """
    result = {}
    with open(csv_path, newline='', encoding='utf-8-sig') as f:
        reader = csv.DictReader(f, delimiter=';')
        # 构建字段映射：原始字段名 -> 清洗后的字段名（小写，去除括号内内容）
        mapping = {}
        for field in reader.fieldnames:
            clean_field = re.sub(r'\s*\([^)]*\)', '', field).strip().lower()
            mapping[field] = clean_field

        # 需要排除的字段（均为清洗后的名称）
        exclude = {"device", "plus form",
                   "diagnosis code", "id", "serie number"}

        for row in reader:
            # 构建 key，假设 CSV 中 ID 列为 "ID"，SERIE NUMBER 列为 "SERIE NUMBER (S)"
            id_val = row.get("ID", "unknown_id").strip()
            serie_val = row.get("SERIE NUMBER (S)", "").strip()
            key = f"{id_val}_{serie_val}"

            # 生成 diagnosis text
            dg_raw = row.get("DIAGNOSIS CODE (DG)", "").strip()
            pf_raw = row.get("PLUS FORM (PF)", "").strip()
            diag_text = diagnosis_mapping.get(dg_raw, dg_raw)
            plus_text = plus_mapping.get(pf_raw, "")
            if plus_text:
                full_text = f"{diag_text}, {plus_text}"
            else:
                full_text = diag_text

            # 构建 patient info 字典，排除指定字段
            info = {}
            for orig_key, value in row.items():
                clean_key = mapping.get(orig_key, orig_key).strip()
                if clean_key in exclude:
                    continue
                # 对 SEX 字段进行处理
                if clean_key == "sex":
                    v = value.strip()
                    if v.upper() == "F":
                        info[clean_key] = "Female"
                    elif v.upper() == "M":
                        info[clean_key] = "Male"
                    else:
                        info[clean_key] = v
                else:
                    info[clean_key] = value.strip()
            result[key] = {"patient info": info, "text": full_text}
    return result


def gather_data(data_path, tar_path, prefix='ridorop_'):
    """
    对 Infant Retinal Database 数据集进行预处理：
      - 数据集根目录下要求存在 images 文件夹和 infant_retinal_database_info.csv 文件；
      - images 文件夹中，目录结构为：images/<patient id>/<serie number>/ ；
      - 对于每个系列，生成 key = "{id}_{serie number}"，从 CSV 中获取该系列的患者信息和诊断文本；
      - 遍历每个系列下的所有图像，对每张图像调用 crop_resize_save 完成裁剪与 resize，
        新图像名称格式为 {prefix}_{patient id}_{serie number}_{原始文件名}.png（保留原始文件名，仅修改后缀并加上前缀）；
      - 将每张图像对应的标注信息写入 data_dict，其中包括 patient info 和 diagnosis text；
      - 最终将 data_dict 保存到 tar_path/annotations.json 中。

    参数：
        data_path (str): 数据集根目录，要求其中包含 images 文件夹和 infant_retinal_database_info.csv 文件；
        tar_path (str): 预处理后数据存放目录；
        prefix (str): 处理后图片名称的前缀，默认为 "ridorop_"

    返回：
        dict: 标注信息字典。
    """
    os.makedirs(tar_path, exist_ok=True)
    images_out_dir = os.path.join(tar_path, 'images')
    os.makedirs(images_out_dir, exist_ok=True)

    src_image_dir = os.path.join(data_path, 'images')
    info_path = os.path.join(data_path, 'infant_retinal_database_info.csv')
    info_dict = parse_info(info_path)

    data_dict = {}

    id_list = sorted(os.listdir(src_image_dir))
    total_id_num = len(id_list)
    # 遍历 src_image_dir 下每个患者文件夹，假设文件夹名称为数字
    for id_str in tqdm(id_list, total=total_id_num, desc='id process', unit='id'):
        patient_folder = os.path.join(src_image_dir, id_str)
        if not os.path.isdir(patient_folder):
            continue
        try:
            id_num = int(id_str)
        except ValueError:
            print(f"Warning: {id_str} is not a valid patient id.")
            continue
        # 遍历患者文件夹中的每个系列文件夹
        for serie_str in sorted(os.listdir(patient_folder)):
            series_folder = os.path.join(patient_folder, serie_str)
            if not os.path.isdir(series_folder):
                continue
            try:
                s_num = int(serie_str)
            except ValueError:
                print(
                    f"Warning: {serie_str} is not a valid serie number in patient {id_str}.")
                continue
            k = f"{str(id_num)}_{str(s_num)}"
            if k not in info_dict:
                print(f"Warning: Key {k} not found in CSV info.")
                continue
            patient_info = info_dict[k].get('patient info', {})
            diag_text = info_dict[k].get('text', "")
            # 遍历该系列下的所有图像文件
            for image_ori_name in sorted(os.listdir(series_folder)):
                if not image_ori_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.ppm')):
                    continue
                base_name = os.path.splitext(image_ori_name)[0]
                new_image_name = f"{prefix}_{id_str}_{serie_str}_{base_name}.png"
                dest_image_path = os.path.join(images_out_dir, new_image_name)
                # 保证目标目录存在
                os.makedirs(images_out_dir, exist_ok=True)
                crop_info = crop_resize_save(
                    image_path=os.path.join(series_folder, image_ori_name),
                    save_path=dest_image_path,
                    resize=(224, 224),
                    crop_threshold=25
                )
                data_dict[new_image_name] = {
                    "image_name": new_image_name,
                    "image_path": os.path.join("images", new_image_name),
                    'original_path': os.path.relpath(os.path.join(series_folder, image_ori_name), data_path),
                    'crop_info': crop_info,
                    "patient_info": patient_info,
                    "diagnosis": {
                        "text": diag_text
                    }
                }

    # 保存 data_dict 到 tar_path/annotations.json
    os.makedirs(tar_path, exist_ok=True)
    annotations_path = os.path.join(tar_path, "annotations.json")
    with open(annotations_path, "w", encoding="utf-8") as f:
        json.dump(data_dict, f, indent=4)

    return data_dict


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(
        description="Infant Retinal Database 数据集预处理")
    parser.add_argument("--data_path", type=str, default="../Dataset/infant_retinal_database",
                        help="数据集根目录，要求其中包含 images 文件夹和 infant_retinal_database_info.csv 文件")
    parser.add_argument("--tar_path", type=str, default="../Dataset/processed224_infant_retinal_database",
                        help="预处理后数据存放目录")
    parser.add_argument("--prefix", type=str, default="ridorop_",
                        help="处理后图片名称的前缀，默认为 'ridorop_'")
    args = parser.parse_args()

    annotations = gather_data(
        args.data_path, args.tar_path, prefix=args.prefix)
    print("WARNING: 本数据集的标注是按照一次诊断的标注，因此，虽然部分图片被标注为了阳性，但不一定能在该图片观察到，可能在同一次随访的拍摄的其他图片中观察病灶")
    print("Preprocessing completed. Annotations saved.")
